use axum::{
    Router,
    response::sse::{Event, Sse},
    routing::get,
};
use futures_util::stream;
use std::{convert::Infallible, sync::Arc, time::Duration};
use tauri::State;
use tauri_axum::{LocalRequest, RequestResult, StreamChunk, StreamManager, handle_request};
use tokio::sync::Mutex;

struct AppState {
    router: Arc<Mutex<Router>>,
    stream_manager: Arc<Mutex<StreamManager>>,
}

#[tauri::command]
async fn handle_app_request(
    state: State<'_, AppState>,
    local_request: LocalRequest,
) -> Result<RequestResult, String> {
    let mut router = state.router.lock().await;
    let mut stream_manager = state.stream_manager.lock().await;

    let result = handle_request(local_request, &mut router, &mut stream_manager).await;
    Ok(result)
}

#[tauri::command]
async fn get_stream_chunk(
    state: State<'_, AppState>,
    stream_id: String,
) -> Result<Option<StreamChunk>, String> {
    let mut stream_manager = state.stream_manager.lock().await;
    let chunk = stream_manager.get_next_chunk(&stream_id).await;
    Ok(chunk)
}

#[tauri::command]
async fn close_stream(state: State<'_, AppState>, stream_id: String) -> Result<(), String> {
    let mut stream_manager = state.stream_manager.lock().await;
    stream_manager.close_stream(&stream_id);
    Ok(())
}

fn create_app_router() -> Router {
    Router::new()
        .route("/", get(|| async { "Hello, World!" }))
        .route("/stream", get(stream_handler))
        .route("/events", get(event_stream_handler))
}

async fn stream_handler() -> Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>> {
    let stream = stream::iter(vec![
        Ok(Event::default().data("First chunk")),
        Ok(Event::default().data("Second chunk")),
        Ok(Event::default().data("Third chunk")),
    ]);

    Sse::new(stream)
}

async fn event_stream_handler() -> Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>
{
    let stream = async_stream::stream! {
        for i in 0..10 {
            tokio::time::sleep(Duration::from_millis(500)).await;
            yield Ok(Event::default().data(format!("Event {}", i)));
        }
    };

    Sse::new(stream)
}

fn main() {
    let app_router = create_app_router();
    let app_state = AppState {
        router: Arc::new(Mutex::new(app_router)),
        stream_manager: Arc::new(Mutex::new(StreamManager::new())),
    };

    tauri::Builder::default()
        .manage(app_state)
        .invoke_handler(tauri::generate_handler![
            handle_app_request,
            get_stream_chunk,
            close_stream
        ])
        .setup(|_app| Ok(()))
        .expect("error while running tauri application");
}

